#!/usr/bin/env python
#
# Generic scheduler components, including a base Machine (scheduler target)
# and a base Scheduler that greedily tries to run jobs as they arrive (if
# possible) or as other jobs finish.
#
# Author:   Mike Murphy (mamurph@cs.clemson.edu)
# Revision: 23 September 2009
#
#   Copyright 2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.


'''
Base components for simulating job scheduling systems. In particular, this
module provides the base Scheduler, which implements a trivial first-come,
first-served scheduling algorithm. Scheduler is designed to be subclassed
to create other schedulers (see the condor_sim module).

SimVOC schedulers perform several tasks in addition to job scheduling. In
particular, schedulers are self-aggregating, meaning that internal statistics
are maintained both for the scheduler pool as a whole and for each Virtual
Organization serviced by the scheduler. Schedulers also provide a resource
broker capability, by which named objects (providing simulated services, such
as storage brokering) can be provided directly to jobs.
'''

from core_common import SC_CLAIMED, SC_FREE, SC_SHUTDOWN, SM_WAITING, \
                        SM_RUNNING, SM_COMPLETED, SM_ERRORS, SM_TOTAL, SM_FREE
from results import SchedResult
#


class Scheduler(object):
   '''
   Generic job scheduler, which starts jobs instantly if nodes are available.
   Otherwise, jobs are queued until running jobs complete, at which time new
   jobs are immediately started. The scheduling discipline is FIFO.
   '''
   def __init__(self, simobj, name='Scheduler', tag='', **kwds):
      '''
      @param simobj: Reference to the simulator object
      @param name: Name of the scheduler
      @type name: str
      @param tag: Optional scheduler tag
      @type tag: str
      @param kwds: Scheduler extra data
      '''
      self.queue = []
      self.running_list = []
      self.targets = []
      self.free_targets = []
      self.metrics = [0, 0, 0, 0, 0, 0]   # wait, run, comp, err, tot, free
      self.sim = simobj
      self.name = name
      self.tag = tag
      self.sched_id = name + '@' + tag + '<' + str(id(self)) + '>'
      self.extra_data = kwds
      self.vo_counts = {}   # vo: [waiting, running, completed, errors]
      self.broker_map = {}
      self.pending_event = None
      #
      # Hack to allow the watchdog to be paused (for simulation performance)
      self.wd_callback = None
   #
   def debug(self, message):
      '''
      Writes debugging messages back to the simulator debugging interface.
      Useful for developing subclasses.

      @param message: Message to be written
      @type message: str
      '''
      self.sim.debug('%s@%s: %s' % (self.name, self.tag, message))
   #

   # Scheduling methods

   def claim_slot(self, machine):
      '''
      Attempts to claim a free machine (CPU core) for the purpose of starting
      a job on that core. If the machine can be claimed, it is removed from
      the list of free scheduling targets.

      @param machine: Machine object whose claim() method will be called

      @rtype: bool
      @return: True iff successful
      '''
      chk = machine.claim()
      if chk:
         self.adjust_metric(SM_FREE, -1)
	 self.free_targets.remove(machine)
      return chk
   #
   def free_slot(self, machine):
      '''
      Marks a machine (CPU core) as free and return it to the list of free
      scheduling targets.

      @param machine: Machine object to mark as free
      '''
      machine.free()
      if machine in self.targets:
         self.adjust_metric(SM_FREE, 1)
	 self.free_targets.append(machine)
   #
   def pick_slot(self, job):
      '''
      Selects a free machine from the free scheduling target list. Since a
      machine could mark itself as claimed, or another scheduler could claim
      it, this method may not find a free machine.

      @param job: Job object. Unused in this base implementation, but may be
                  used by subclasses for purposes such as requirements
                  matching.

      @return: Machine object in the free scheduling targets list if one is
      available (unclaimed). If there are no unclaimed targets, or the free
      list is empty, this method will return None.
      '''
      match = None
      # A machine could presumably mark itself claimed, so we actually do need
      # a loop here. When only 1 scheduler is used, this operation should be
      # O(1) in most cases.
      for m in self.free_targets:
         if not m.claimed:
	    match = m
	    break
      return match
   #
   def schedule(self, time, simobj):
      '''
      Core FIFO/FCFS scheduling algorithm. This method should be called
      whenever jobs complete and others are waiting, or when jobs arrive and
      there are free slots available. Subclasses should generally either
      override this method or bypass this method by overriding other methods
      that call this method.

      @param time: Current simulated time
      @type time: number
      @param simobj: Reference to the simulator object

      @return: result collection (1-tuple containing a list of results) for
      simulator output
      '''
      r = []
      progress_made = True
      while (self.get_metric(SM_FREE) > 0) and \
            (self.get_metric(SM_WAITING) > 0) and progress_made:
         job = self.queue[-1]
         node = self.pick_slot(job)
         if node is not None:
	    if self.claim_slot(node):
	       r.append(self.start_job(job, node))
            else:
               progress_made = False
         else:
            progress_made = False
         #
      #
      if not progress_made:
         # Nodes are being used by another scheduler: wait a while
         self.check_schedule(300)
      r.append(self.get_update())
      return (r,)
   #
   def event_pending(self):
      '''
      @rtype: bool
      @return: True iff a scheduling event is pending
      '''
      wait = False
      if self.pending_event is not None:
         wait = not self.pending_event.activated
      return wait
   #
   def check_schedule(self, delay=0):
      '''
      Conditionally adds a scheduling event (a simulator-induced call to the
      schedule() method) if such an event is not already pending AND there
      are jobs waiting in the queue.

      @param delay: Optional scheduling delay to add
      @type delay: number

      @rtype: bool
      @return: True iff a scheduling event was added
      '''
      result = False
      if not self.event_pending():
         if self.get_metric(SM_WAITING) > 0:
	    when = self.sim.get_time() + delay
	    self.add_scheduler_event(when)
	    result = True
      return result
   #
   def add_scheduler_event(self, time):
      '''
      Adds a scheduler event for a time not in the past. An Exception will be
      raised if a scheduler event is already pending, or if the time is in the
      past.

      @param time: Time at which the scheduler should run
      @type time: number
      '''
      if self.event_pending():
         raise Exception('Scheduler event already pending')
      self.pending_event = self.sim.new_periodic_event(time, self.schedule)
   #

   # Job management methods

   def enqueue_job(self, job):
      '''
      Adds a job to the scheduler queue. If a watchdog callback function has
      been set, then the callback will be activated.

      @param job: Job to add to the queue

      @return: a result collection or single result for simulator output
      '''
      time = self.sim.get_time()
      r = job.submit(time, self)
      self.queue.insert(0, job)
      self.adjust_metric(SM_WAITING, 1)
      self.adjust_vo(job.vo, SM_WAITING, 1)
      if self.wd_callback is not None:
         self.wd_callback()
      return r
   #
   def submit_job(self, job):
      '''
      Receives a job into the scheduler. This method should be called either
      directly or as a result of a submit event executing, whenever a job needs
      to be added to the scheduler queue.

      @param job: Job to be submitted

      @return: a result collection (1-tuple containing a list)
      '''
      r = []
      r.append(self.enqueue_job(job))
      time = self.sim.get_time()
      self.check_schedule()
      return (r,)
   #
   def submit_trigger(self, job):
      '''
      Method to be called by a simulator event whenever a job is scheduled
      for submission. This method simply calls the submit_job() method and
      exists simply to allow subclasses to override it.

      @param job: Job to be submitted

      @return: a result collection
      '''
      return self.submit_job(job)
   #
   def add_submit_event(self, time, job):
      '''
      Schedules an event to submit a job at a time not in the past.

      @param time: Time at which job should be submitted
      @type time: number
      @param job: Job to submit
      '''
      self.sim.new_external_event(time, self.submit_trigger, job)
   #
   def start_job(self, job, machine):
      '''
      Starts a job on a machine. The machine must already have been claimed
      using the claim_slot() method.

      @param job: Job to be started
      @param machine: Claimed machine on which job is to be run

      @rtype: L{results.JobStartResult}
      '''
      time = self.sim.get_time()
      self.queue.remove(job)
      self.running_list.append(job)
      self.adjust_metric(SM_WAITING, -1)
      self.adjust_metric(SM_RUNNING, 1)
      vo = job.vo
      self.adjust_vo(vo, SM_WAITING, -1)
      self.adjust_vo(vo, SM_RUNNING, 1)
      return job.start(time, self.sim, machine)
   #
   def mark_complete(self, machine, job):
      '''
      Marks a job as complete, adjusting the metrics and freeing the
      associated slot.

      @param machine: Machine on which job was running
      @param job: Job to be marked complete
      '''
      self.running_list.remove(job)
      self.adjust_metric(SM_RUNNING, -1)
      vo = job.vo
      self.adjust_vo(vo, SM_RUNNING, -1)
      #
      if job.error:
         self.adjust_metric(SM_ERRORS, 1)
         self.adjust_vo(vo, SM_ERRORS, 1)
      #
      self.adjust_metric(SM_COMPLETED, 1)
      self.adjust_vo(vo, SM_COMPLETED, 1)
      self.free_slot(machine)
   #
   def job_done(self, machine, job):
      '''
      Callback method used by a Job to notify the scheduler that the job
      is finishing.

      @param machine: Machine on which the job was running
      @param job: Job that is finishing

      @return: scheduler status result collection
      '''
      self.mark_complete(machine, job)
      self.check_schedule()
      return self.get_update()
   #

   def cancel_job(self, job):
      '''
      Cancels a job that has NOT YET STARTED. Exceptions will occur if this
      method is run on a job that has already been scheduled. It is recommended
      that client code call the kill_job() method instead.

      @param job: Job to be canceled

      @return: scheduler status result collection
      '''
      self.queue.remove(job)
      self.adjust_metric(SM_WAITING, -1)
      self.adjust_vo(job.vo, SM_WAITING, -1)
      return self.get_update()
   #
   def kill_job(self, job):
      '''
      Immediately kills a job, whether the job is running or not. If the job
      has not yet been scheduled, it will be canceled. The job must not yet
      have completed (i.e. must be either in the queue or on the running
      list), or an Exception will be raised.

      @param job: Job to kill
      '''
      if (job not in self.queue) and (job not in self.running_list):
         raise Exception('Job has completed or is not known to the scheduler')
      cur_time = self.sim.get_time()
      job.kill(cur_time, self.sim)
      if not job.running:
         self.cancel_job(job)
   #
   def kill_all(self, field_name, value):
      '''
      Kills all jobs for which the field name matches a specified value. To
      kill all jobs known to the scheduler, pass None to both arguments. The
      special field name "job_id" matches the job name, while "vo" matches the
      Virtual Organization with which the job is associated.

      @param field_name: Name of job extra_data field or above special name
      @type field_name: str
      @param value: Value of the field to match (test uses ==)

      @rtype: list
      @return: list of jobs that have been killed
      '''
      all_jobs = []
      all_jobs.extend(self.running_list)
      all_jobs.extend(self.queue)
      killed = []
      for job in all_jobs:
         if (field_name == 'job_id') and (job.name == value):
            killed.append(job)
         elif (field_name == 'vo') and (job.vo == value):
            killed.append(job)
         elif (field_name in job.extra_data):
            if (job.extra_data[field_name] == value):
               killed.append(job)
         elif (field_name is None) and (value is None):
            killed.append(job)
      #
      cur_time = self.sim.get_time()
      for job in killed:
         self.kill_job(job)
      return killed
   #

   # Pool management methods

   def add_target(self, machine):
      '''
      Adds a target machine to the scheduler pool. A single Machine instance
      can occur in the pool at most once. Adding a machine already present in
      the pool will result in an Exception being raised.

      @param machine: Target machine to add
      '''
      if machine not in self.targets:
         self.targets.append(machine)
         self.adjust_metric(SM_TOTAL, 1)
	 if not machine.claimed:
            self.adjust_metric(SM_FREE, 1)
	    self.free_targets.append(machine)
	 self.check_schedule()
      else:
         raise Exception('Duplicate machine')
   #
   def add_targets(self, machine_list):
      '''
      Adds a list of scheduling targets to the pool. Each target in the list
      must be unique, and no target may already be present in the pool,
      otherwise an Exception will occur.

      @param machine_list: List or sequence of target machines to add
      '''
      for machine in machine_list:
         self.add_target(machine)
   #
   def remove_target(self, machine):
      '''
      Removes a target machine from the scheduler pool. If the target machine
      is not in the pool, an Exception will be raised.

      @param machine: Target machine to remove from the pool
      '''
      self.targets.remove(machine)
      self.adjust_metric(SM_TOTAL, -1)
      if not machine.claimed:
         self.adjust_metric(SM_FREE, -1)
	 self.free_targets.remove(machine)
   #
   def remove_targets(self, machine_list):
      '''
      Removes a list of target machines from the scheduler pool. Each target
      machine must be unique in the list and must already be present in the
      pool, or an Exception will be raised.

      @param machine_list: List or sequence of target machines to remove
      '''
      for machine in machine_list:
         self.remove_target(machine)
   #

   # Self-aggregation and statistical routines

   def adjust_metric(self, field, adj):
      '''
      Changes an aggregation metric field. This aggregation is used to count
      the numbers of jobs waiting, running, and completed, as well as to
      total errors, job count, and free nodes.

      @param field: Field number (use SM_ constant from L{core_common})
      @type field: int
      @param adj: Adjustment (positive or negative offset)
      @type adj: int
      '''
      cur = self.metrics[field]
      self.set_metric(field, cur + adj)
   #
   def get_metric(self, field):
      '''
      Returns the value of the metric for a given field. Field constants have
      the SM_ prefix and may be found in the core_common module.

      @param field: Field number (use SM_ constant)
      @type field: int
      '''
      return self.metrics[field]
   #
   def set_metric(self, field, qty):
      '''
      Sets the value of an aggregation metric. Field constants, with the SM_
      prefix, may be found in core_common.

      @param field: Field number (use SM_ constant)
      @type field: int
      @param qty: Value to which the field is to be set
      @type qty: int
      '''
      self.metrics[field] = qty
   #
   def adjust_vo(self, vo, field, adj):
      '''
      Adjust the aggregation metrics for a single VO. Field constants, with
      the SM_ prefix, may be found in the core_common module. Note that
      SM_TOTAL and SM_FREE are not tracked for VO data.

      @param vo: Virtual Organization name
      @type vo: str
      @param field: Field number (use SM_ constant)
      @type field: int
      @param adj: Adjustment factor (positive or negative offset)
      @type adj: int
      '''
      cur = 0
      if vo in self.vo_counts:
         cur = self.vo_counts[vo][field]
      self.set_vo(vo, field, cur + adj)
   #
   def get_vo(self, vo, field):
      '''
      Returns an aggregation metric for a single VO. Field constants have the
      prefix SM_ and may be found in core_common.

      @param vo: Virtual Organization name
      @type vo: str
      @param field: Field number (use SM_ constant)
      @type field: int
      '''
      count = 0
      if vo in self.vo_counts:
         count = self.vo_counts[vo][field]
      return count
   #
   def set_vo(self, vo, field, qty):
      '''
      Sets the value of an aggregation metric for a specific VO. If the VO
      has not yet been registered, it will be added to the aggregation
      structure. Field names, with the prefix SM_, may be found in the
      core_common module.

      @param vo: Virtual Organization name
      @type vo: str
      @param field: Field number (use SM_ constant)
      @type field: int
      @param qty: Value to which the specified field should be set
      @type qty: int
      '''
      if vo not in self.vo_counts:
         self.vo_counts[vo] = [0, 0, 0, 0]
      #
      self.vo_counts[vo][field] = qty
   #
   def get_update(self):
      '''
      Obtains information about the current state of the scheduler.

      @rtype: L{results.SchedResult}
      '''
      ts = self.sim.get_timestamp()
      r = SchedResult(ts, self.sched_id, self.name, self.tag, \
                      self.metrics, self.vo_counts)
      return r
   #
   def set_watchdog_callback(self, wdfn):
      '''
      Sets a callback function to be called whenever a job is added to the
      scheduler queue. This callback function is intended for use with VOC
      simulations that use a Watchdog component to monitor the queue size.

      @param wdfn: Function to be called
      '''
      self.wd_callback = wdfn
   #
   def get_count(self, field, vo=None):
      '''
      Returns the value of a metric for the scheduler aggregation or for a
      single Virtual Organization. Field constants have the prefix SM_ and
      may be found in the core_common module.

      @param field: Field number (use SM_ constant)
      @type field: int
      @param vo: VO name (str) or None (for scheduler aggregation)
      '''
      count = 0
      if vo is not None:
         count = self.get_vo(vo, field)
      else:
         count = self.get_metric(field)
      return count
   #
   def get_waiting_count(self, vo=None):
      '''
      @param vo: VO name (str) or None (for whole scheduler data)

      @rtype: int
      @return: number of jobs waiting in the queue for a single VO or for
      the entire scheduler
      '''
      return self.get_count(SM_WAITING, vo)
   #
   def get_running_count(self, vo=None):
      '''
      @param vo: VO name (str) or None (for whole scheduler data)

      @rtype: int
      @return: number of jobs currently running for a single VO or for
      the entire scheduler
      '''
      return self.get_count(SM_RUNNING, vo)
   #
   def get_error_count(self, vo=None):
      '''
      @param vo: VO name (str) or None (for whole scheduler data)

      @rtype: int
      @return: number of job errors observed for a single VO or for
      the entire scheduler.
      '''
      return self.get_count(SM_ERRORS, vo)
   #
   def get_total_queue_size(self, vo=None):
      '''
      @param vo: VO name (str) or None (for whole scheduler data)

      @rtype: int
      @return: total queue size (waiting+running) for a single VO or for
      the entire scheduler.
      '''
      return self.get_waiting_count(vo) + self.get_running_count(vo)
   #
   def get_completed_count(self, vo=None):
      '''
      @param vo: VO name (str) or None (for whole scheduler data)

      @rtype: int
      @return: number of jobs that have completed for a single VO or for
      the entire scheduler.
      '''
      return self.get_count(SM_COMPLETED, vo)
   #
   def is_target(self, machine):
      '''
      @param machine: Machine to test for membership in the pool

      @rtype: bool
      @return: True iff machine is part of the scheduler pool
      '''
      return machine in self.targets
   #
   def get_free_count(self):
      '''
      @rtype: int
      @return: total number of free scheduler slots
      '''
      return self.get_count(SM_FREE)
   #
   def get_cpu_count(self):
      '''
      @rtype: int
      @return: total number of scheduler slots
      '''
      return self.get_count(SM_TOTAL)
   #

   # Resource broker capability

   def add_resource(self, service_name, obj):
      '''
      Adds a resource to the resource broker map.

      @param service_name: Name of the resource
      @type service_name: str
      @param obj: Resource object
      '''
      self.broker_map[service_name] = obj
   #
   def remove_resource(self, service_name):
      '''
      Removes a resource from the resource broker map. The resource must be
      present in the map, or an Exception will occur.

      @param service_name: Name of resource to remove
      @type service_name: str
      '''
      del self.broker_map[service_name]
   #
   def get_resource(self, service_name):
      '''
      @param service_name: Name of the resource to retrieve
      @type service_name: str

      @return: a reference to a resource object, if the resource is available
      from the scheduler. If the resource is not available, this method will
      return None.
      '''
      obj = None
      if service_name in self.broker_map:
         obj = self.broker_map[service_name]
      return obj
   #
#


if __name__ == '__main__':
   from cluster_sim import Cluster
   from job_model import Job
   from kernel import Simulator

   sim = Simulator()
   c = Cluster('test', sim)
   c.grow(10)
   s = Scheduler(sim)
   c.add_scheduler('Scheduler', s)
   s.add_targets(c.machines)

   for x in xrange(0,20):
      job = Job(20 + (x % 4), name='job%d' % x, prob_err=0.2)
      s.add_submit_event(5, job)
   #

   sim.run()
#
